package com.cctc.rds.raft.core.util;

import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.BooleanSupplier;

public class CountDownWaiter {
    private int count;
    private int[] metrics;
    private Lock lock = new ReentrantLock();
    private Condition condition = lock.newCondition();


    public CountDownWaiter(int count) {
        this.count = count;
        this.metrics = new int[3];
    }

    public void countDown() {
        countDown(0);
    }

    public void countDownSuccess() {
        countDown(1);
    }

    public void countDownException() {
        countDown(2);
    }

    public void await() {
        await(() -> metrics[0] >= count);
    }

    public void await(int minSuccess) {
        await(() -> metrics[0] >= count || metrics[1] >= minSuccess);
    }

    public int getCount() {
        return getCount(0);
    }

    public int getCountSuccess() {
        return getCount(1);
    }

    public int getCountException() {
        return getCount(2);
    }

    private void countDown(int index) {
        lock.lock();
        try {
            metrics[index]++;
            condition.signal();
        } finally {
            lock.unlock();
        }
    }

    private void await(BooleanSupplier supplier) {
        lock.lock();
        try {
            while (!Thread.interrupted()) {
                if (supplier.getAsBoolean()) return;
                condition.await();
            }
        } catch (InterruptedException ignored) {
        } finally {
            lock.unlock();
        }
    }

    private int getCount(int index) {
        lock.lock();
        try {
            return metrics[index];
        } finally {
            lock.unlock();
        }
    }

}
